import argparse
import ctypes

import numpy.ctypeslib as npct
from torchsummary import summary

from datasets import *
from detect import *
from models.yolo import Model
from utils.utils import *

libLoad = ctypes.cdll.LoadLibrary
sharelib = libLoad("/home/st/ubuntu_data/test_data/slam_det_tra_pre_ws/devel/lib/libpython_share_memory.so") #need change 

img_flag_ = ctypes.c_int(100)
p_img_flag_ = ctypes.pointer(img_flag_)
img_ = np.zeros(shape=(480,640,3), dtype=np.uint8)
sharelib.get_img.argtypes = [
    npct.ndpointer(dtype=np.uint8, ndim=3, shape =img_.shape),
]

def get_img():
    sharelib.get_img(img_)
    return img_
 
def get_img_flag():
    sharelib.get_img_flag(p_img_flag_)
    return img_flag_.value

def send_img_pre_data(mea_data):
    img_pre = []
    img_pre = np.array(img_pre)
    if mea_data is not None:
        img_pre = np.zeros([mea_data.shape[0],3])
        img_pre[:,0] =  mea_data[:, 5] 
        img_pre[:,1]  =   (mea_data[:, 0] + mea_data[:,2])*0.5
        img_pre[:,2]  =  (mea_data[:, 1] + mea_data[:,3])*0.5
    img_pre = np.array( img_pre,dtype=np.float32)
    if not img_pre.flags['C_CONTIGUOUS']:
        img_pre = np.ascontiguous(img_pre, dtype=img_pre.dtype)  # 如果不是C连续的内存，必须强制转换
    p_img_pre = ctypes.cast(img_pre.ctypes.data, ctypes.POINTER(ctypes.c_float))   #转换为ctypes，这里转换后的可以直接利用ctypes转换为c语言中的int*，然后在c中使用
    sharelib.send_img_pre_data(p_img_pre, img_pre.shape[0])

class yolov5:
    def __init__(self, model_file, model_cfg, imagesz = 640):
        self.weights = model_file
        self.imgsz = imagesz
        self.device = torch_utils.select_device('0') 
        self.augment = False
        self.conf_thres = 0.73
        self.iou_thres = 0.5
        self.classes = None
        self.agnostic_nms = False
        self.model = Model(model_cfg).to(self.device)
        # self.model = torch.load(self.weights, map_location=self.device)['model']
        self.model.load_state_dict(torch.load(self.weights))
        self.model.to(self.device).eval()
        img = torch.zeros((1, 3, self.imgsz, self.imgsz), device=self.device)  # init img
        _ = self.model(img.float()) if self.device.type != 'cpu' else None  # run once

    def detect(self, image_raw):
        image_det = image_raw.copy()
        image_det = letterbox(image_det, new_shape=self.imgsz)[0]
        image_det = image_det[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
        image_det = np.ascontiguousarray(image_det)

        image_det = torch.from_numpy(image_det).to(self.device)
        image_det = image_det.float()  # uint8 to fp16/32
        image_det /= 255.0  # 0 - 255 to 0.0 - 1.0
        if image_det.ndimension() == 3:
            image_det = image_det.unsqueeze(0)
        
        # Inference
        pred = self.model(image_det, augment=self.augment)[0]

        # Apply NMS
        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres,\
            fast=True, classes=self.classes, agnostic=self.agnostic_nms)
        det = pred[0]
        # gn = torch.tensor(image_raw.shape)[[1, 0, 1, 0]]  #  normalization gain whwh
        if det is not None and len(det):
            # Rescale boxes from img_size to im0 size
            det[:, :4] = scale_coords(image_det.shape[2:], det[:, :4], image_raw.shape).round()  # x1y1x2y2,  conf, cls
            
        return det

class Image_detection: 
    def __init__(self ):
        self.image_detector = yolov5("/home/st/ubuntu_data/NRSL/object_detection/history/yolov5_learning/weights/best_weight.pt", \
            "/home/st/ubuntu_data/NRSL/object_detection/history/yolov5_learning/models/score/my_model_setting0626.yaml", 640)

    def img_detection(self, cv_image):
        det = self.image_detector.detect(cv_image)
        if det is None or len(det) == 0:
            return

        det = det.detach().cpu().numpy()
        return det

if __name__ == "__main__":    
    image_det=Image_detection()
    print("yolov5 launch!")
    while True:
        time.sleep(0.001)
        # print("img flag ", get_img_flag())
        if get_img_flag()==1 :
            t0 = time.time()
            img = get_img()
            send_img_pre_data(image_det.img_detection(img))
            torch.cuda.empty_cache()
            t1= time.time() - t0
            print("time ", t1)

